import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
from xgboost import XGBRegressor
from statsmodels.tsa.seasonal import seasonal_decompose  # 导入 seasonal_decompose
import warnings
warnings.filterwarnings('ignore')
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 设置页面标题和布局
st.set_page_config(page_title="多城市空气质量分析与预测系统", layout="wide")

# 标题
st.title("多城市空气质量分析与预测系统 (2020-2024)")

# 加载数据
@st.cache_data
def load_data():
    try:
        data = pd.read_csv('processed_air_quality.csv')
        data['date'] = pd.to_datetime(data['date'], format='%Y-%m-%d')
        data['year'] = data['date'].dt.year
        data['month'] = data['date'].dt.month
        data['day'] = data['date'].dt.day
        data['weekday'] = data['date'].dt.weekday
        daily_data = data.groupby(['date', 'year', 'month', 'day', 'weekday']).mean().reset_index()
        return data, daily_data
    except FileNotFoundError:
        st.error("数据文件未找到，请检查文件名和路径是否正确。")
        return None, None

data, daily_data = load_data()

# 检查数据是否加载成功
if data is None:
    st.stop()

# 侧边栏 - 城市选择和分析选项
st.sidebar.header("分析选项")
selected_cities = st.sidebar.multiselect(
    "选择城市", 
    ['北京', '上海', '广州'], 
    default=['北京']
)

analysis_type = st.sidebar.selectbox(
    "选择分析类型",
    ["数据概览", "时间序列分析", "污染物相关性", "AQI预测"]
)

# 主内容区域
if analysis_type == "数据概览":
    st.header("数据概览")
    
    # 显示原始数据
    if st.checkbox("显示原始数据"):
        st.write(data.head())
    
    # AQI分布
    st.subheader("AQI分布")
    fig, ax = plt.subplots(figsize=(10, 6))
    for city in selected_cities:
        sns.kdeplot(data[f'AQI_{city}'], label=city, ax=ax)
    ax.set_title("AQI分布密度图")
    ax.set_xlabel("AQI值")
    ax.set_ylabel("密度")
    ax.legend()
    st.pyplot(fig)
    
    # 年度AQI变化
    st.subheader("年度AQI变化")
    yearly_aqi = daily_data.groupby(['year'])[[f'AQI_{city}' for city in selected_cities]].mean()
    st.line_chart(yearly_aqi)
    
    # 污染物对比
    st.subheader("主要污染物对比")
    pollutants = ['PM2.5', 'PM10', 'SO2', 'NO2', 'O3', 'CO']
    selected_pollutant = st.selectbox("选择污染物", pollutants)
    
    fig, ax = plt.subplots(figsize=(10, 6))
    for city in selected_cities:
        sns.lineplot(
            x='date', 
            y=f'{selected_pollutant}_{city}', 
            data=daily_data, 
            label=city, 
            ax=ax
        )
    ax.set_title(f"{selected_pollutant}时间变化")
    ax.set_xlabel("日期")
    ax.set_ylabel(f"{selected_pollutant}浓度")
    ax.legend()
    st.pyplot(fig)
    
    # 交互式污染物对比分析
    st.subheader("交互式污染物对比分析")
    
    col1, col2, col3 = st.columns(3)
    with col1:
        selected_pollutants = st.multiselect(
            "选择要对比的污染物",
            ['AQI', 'PM2.5', 'PM10', 'SO2', 'NO2', 'O3', 'CO'],
            default=['AQI', 'PM2.5']
        )
    with col2:
        compare_cities = st.multiselect(
            "选择要对比的城市",
            ['北京', '上海', '广州'],
            default=['北京', '上海']
        )
    with col3:
        min_date = data['date'].min().date()
        max_date = data['date'].max().date()
        date_range = st.date_input(
            "选择时间范围",
            value=[min_date, max_date],
            min_value=min_date,
            max_value=max_date
        )
    
    if len(selected_pollutants) > 0 and len(compare_cities) > 0:
        start_date = pd.to_datetime(date_range[0])
        end_date = pd.to_datetime(date_range[1])
        mask = (daily_data['date'] >= start_date) & (daily_data['date'] <= end_date)
        filtered_data = daily_data.loc[mask]
        
        tabs = st.tabs(selected_pollutants)
        for tab, pollutant in zip(tabs, selected_pollutants):
            with tab:
                fig, ax = plt.subplots(figsize=(12, 5))
                for city in compare_cities:
                    col_name = f'{pollutant}_{city}'
                    if col_name in filtered_data.columns:
                        sns.lineplot(
                            x='date',
                            y=col_name,
                            data=filtered_data,
                            label=city,
                            ax=ax
                        )
                ax.set_title(f"{pollutant}浓度时间序列对比")
                ax.set_xlabel("日期")
                ax.set_ylabel(f"{pollutant}浓度")
                ax.legend()
                ax.grid(True)
                st.pyplot(fig)
                
                st.subheader(f"{pollutant}统计摘要")
                stats_data = []
                for city in compare_cities:
                    col_name = f'{pollutant}_{city}'
                    if col_name in filtered_data.columns:
                        stats = filtered_data[col_name].describe().to_frame(name=city)
                        stats_data.append(stats)
                if stats_data:
                    combined_stats = pd.concat(stats_data, axis=1)
                    st.dataframe(combined_stats.style.format("{:.2f}"))
    else:
        st.warning("请至少选择一个污染物和一个城市进行对比")

elif analysis_type == "时间序列分析":
    st.header("时间序列分析")
    
    selected_city = st.selectbox("选择城市", selected_cities)
    selected_pollutant = st.selectbox(
        "选择指标", 
        ['AQI', 'PM2.5', 'PM10', 'SO2', 'NO2', 'O3', 'CO']
    )
    
    st.subheader(f"{selected_city} {selected_pollutant}时间序列")
    ts_data = daily_data.set_index('date')[[f'{selected_pollutant}_{selected_city}']].dropna()
    st.line_chart(ts_data)
    
    st.subheader("季节性分解")
    min_data_length = 365
    if len(ts_data) > min_data_length:  
        result = seasonal_decompose(ts_data, model='additive', period=365)
        fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1, figsize=(12, 8))
        result.observed.plot(ax=ax1)
        ax1.set_ylabel('Observed')
        result.trend.plot(ax=ax2)
        ax2.set_ylabel('Trend')
        result.seasonal.plot(ax=ax3)
        ax3.set_ylabel('Seasonal')
        result.resid.plot(ax=ax4)
        ax4.set_ylabel('Residual')
        plt.tight_layout()
        st.pyplot(fig)
    else:
        st.warning(f"数据不足，至少需要{min_data_length}个数据点才能进行季节性分解")

elif analysis_type == "污染物相关性":
    st.header("污染物相关性分析")
    
    selected_city = st.selectbox("选择城市", selected_cities)
    
    st.subheader(f"{selected_city} 污染物相关性")
    pollutants = ['AQI', 'PM2.5', 'PM10', 'SO2', 'NO2', 'O3', 'CO']
    corr_data = daily_data[[f'{p}_{selected_city}' for p in pollutants]].corr()
    
    fig, ax = plt.subplots(figsize=(10, 8))
    sns.heatmap(corr_data, annot=True, cmap='coolwarm', ax=ax)
    ax.set_title(f"{selected_city} 污染物相关性热力图")
    st.pyplot(fig)
    
    st.subheader("污染物关系散点图矩阵")
    fig = px.scatter_matrix(
        daily_data,
        dimensions=[f'{p}_{selected_city}' for p in pollutants],
        title=f"{selected_city} 污染物关系散点图矩阵"
    )
    st.plotly_chart(fig)

elif analysis_type == "AQI预测":
    st.header("AQI预测模型")
    
    selected_city = st.selectbox("选择预测城市", selected_cities)
    target_col = f'AQI_{selected_city}'
    feature_cols = [f'{p}_{selected_city}' for p in ['PM2.5', 'PM10', 'SO2', 'NO2', 'O3', 'CO']]
    
    pred_data = daily_data[['date', target_col] + feature_cols].dropna()
    X = pred_data[feature_cols]
    y = pred_data[target_col]
    
    train_size = int(len(pred_data) * 0.8)
    train, test = pred_data.iloc[:train_size], pred_data.iloc[train_size:]
    X_train, X_test = X.iloc[:train_size], X.iloc[train_size:]
    y_train, y_test = y.iloc[:train_size], y.iloc[train_size:]
    
    model_type = st.selectbox(
        "选择预测模型",
        ["线性回归", "随机森林", "XGBoost"]
    )
    
    if model_type in ["线性回归", "随机森林", "XGBoost"]:
        if model_type == "线性回归":
            model = LinearRegression()
        elif model_type == "随机森林":
            n_estimators = st.slider("选择树的数量", 10, 200, 100)
            model = RandomForestRegressor(n_estimators=n_estimators, random_state=42)
        elif model_type == "XGBoost":
            n_estimators = st.slider("选择树的数量", 10, 200, 100)
            max_depth = st.slider("选择树的最大深度", 3, 10, 5)
            learning_rate = st.slider("选择学习率", 0.01, 0.3, 0.1)
            model = XGBRegressor(
                n_estimators=n_estimators,
                max_depth=max_depth,
                learning_rate=learning_rate,
                random_state=42
            )
        
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        
        # 评估指标
        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)
        
        # 预测结果展示
        if 'y_pred' in locals() and len(test) > 0:
            st.subheader("测试集日期范围")
            st.caption(f"{test['date'].min().date()} 至 {test['date'].max().date()}")
            
            st.subheader("单日预测对比")
            col1, col2 = st.columns(2)
            
            with col1:
                predict_date = st.date_input(
                    "选择预测日期",
                    value=test['date'].iloc[0].date(),
                    min_value=test['date'].min().date(),
                    max_value=test['date'].max().date()
                )
            
            try:
                date_mask = (test['date'] == pd.to_datetime(predict_date))
                if date_mask.any():
                    predict_idx = date_mask.idxmax()
                    
                    actual = y_test.loc[predict_idx]
                    predicted = y_pred[predict_idx - test.index[0]]
                    
                    with col2:
                        st.metric("实际AQI", f"{actual:.1f}")
                        st.metric("预测AQI", f"{predicted:.1f}", delta=f"{(predicted-actual):.1f}")
                else:
                    st.warning("选择的日期不在测试集范围内")
                    
            except Exception as e:
                st.error(f"预测出错: {str(e)}")
            
            # 动态预测对比图表
            st.subheader("动态预测对比")
            window_size = st.slider("选择时间窗口大小(天)", 7, 90, 30, key="window_size")
            
            compare_df = test[['date', target_col]].copy()
            compare_df['预测值'] = y_pred
            
            fig = px.line(
                compare_df,
                x='date',
                y=[target_col, '预测值'],
                labels={'value': 'AQI', 'variable': '类型'},
                title=f"实际值与预测值动态对比"
            )
            
            fig.update_xaxes(
                rangeslider_visible=True,
                rangeselector=dict(
                    buttons=list([
                        dict(count=window_size, label=f"{window_size}天", step="day", stepmode="backward"),
                        dict(count=window_size*3, label=f"{window_size*3}天", step="day", stepmode="backward"),
                        dict(step="all")
                    ])
                )
            )
            
            fig.update_traces(
                hovertemplate="日期: %{x|%Y-%m-%d}<br>AQI: %{y:.1f}",
                line=dict(width=2)
            )
            fig.update_layout(
                hovermode="x unified",
                legend_title_text="数据类型"
            )
            
            st.plotly_chart(fig, use_container_width=True)
            
            # 模型评估指标
            st.subheader("模型评估")
            col1, col2 = st.columns(2)
            with col1:
                st.metric("均方误差(MSE)", f"{mse:.2f}")
            with col2:
                st.metric("R²分数", f"{r2:.2f}")

# 添加项目说明
st.sidebar.markdown("""
### 项目说明
本系统对2020-2024年多个城市的空气质量数据进行分析和预测，包含以下功能:
1. 数据概览 - 查看AQI分布和污染物变化
2. 时间序列分析 - 分析趋势和季节性
3. 污染物相关性 - 探索污染物间关系
4. AQI预测 - 使用不同模型预测AQI
""")